package com.btl.service.ai.biz.facade;

import cn.idev.excel.FastExcel;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.baomidou.mybatisplus.core.toolkit.IdWorker;
import com.btl.boot.exception.BizException;
import com.btl.boot.exception.SysException;
import com.btl.service.ai.auth.SecurityContextHolder;
import com.btl.service.ai.biz.constants.DocTypeEnum;
import com.btl.service.ai.biz.constants.SegmentStatusEnum;
import com.btl.service.ai.biz.constants.TaskStatusEnum;
import com.btl.service.ai.biz.entity.KwDocument;
import com.btl.service.ai.biz.entity.KwSegment;
import com.btl.service.ai.biz.service.IKwDocumentService;
import com.btl.service.ai.biz.service.IKwSegmentQuestionService;
import com.btl.service.ai.biz.service.IKwSegmentService;
import com.btl.service.ai.biz.queue.DocumentIngestQueue;
import com.btl.service.ai.biz.queue.DocumentIngestTask;
import com.btl.service.ai.support.input.reply.FaqData;
import com.btl.service.ai.support.input.reply.FaqExcelDataListener;
import com.btl.service.ai.support.rag.embedding.strategy.EmbeddingModelTypeEnum;
import com.btl.service.ai.support.rag.embedding.strategy.EmbeddingStoreIngestorStrategy;
import com.btl.service.ai.support.rag.embedding.strategy.EmbeddingStoreIngestorStrategyFactory;
import com.btl.service.ai.support.input.common.parse.strategy.FileTypeEnum;
import com.btl.service.ai.support.input.common.parse.strategy.ParseStrategy;
import com.btl.service.ai.support.input.common.parse.strategy.ParseStrategyFactory;
import com.btl.service.ai.support.input.common.split.strategy.SplitStrategyFactory;
import com.btl.service.ai.support.input.common.split.strategy.SplitterEnum;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.segment.TextSegment;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;
import org.springframework.transaction.support.TransactionSynchronization;
import org.springframework.transaction.support.TransactionSynchronizationManager;
import org.springframework.web.multipart.MultipartFile;

import java.io.IOException;
import java.time.LocalDateTime;
import java.util.*;

/**
 * 文档处理
 */
@Slf4j
@Service
@RequiredArgsConstructor
public class DocumentFacade {

    private final IKwSegmentService iKwSegmentService;

    private final IKwSegmentQuestionService iKwSegmentQuestionService;

    private final IKwDocumentService iKwDocumentService;

    private final DocumentIngestQueue documentIngestQueue;

    /**
     * 创建文档-通用
     *
     * @param file     文件
     * @param splitter 分割器
     * @return Long
     */
    @Transactional(rollbackFor = Exception.class)
    public Long createCommon(SplitterEnum splitter, MultipartFile file) {
        Long childId = SecurityContextHolder.get()
                .getChildId();
        Long documentId;
        Document document;
        String fileName = file.getOriginalFilename();
        try {
            ParseStrategy parseStrategy = ParseStrategyFactory.getStrategy(FileTypeEnum.fromContentType(file.getContentType()));
            document = parseStrategy.parse(file.getInputStream());
            DocumentSplitter documentSplitter = SplitStrategyFactory.getDocumentSplitter(splitter);
            documentId = IdWorker.getId();
            List<TextSegment> segments = documentSplitter.split(document);
            List<KwSegment> kwSegments = new ArrayList<>();
            int seqId = 1;
            for (TextSegment segment : segments) {
                Long id = IdWorker.getId();
                int codePointCount = segment.text()
                        .codePointCount(0, segment.text()
                                .length());
                KwSegment kwSegment = KwSegment.builder()
                        .childId(childId)
                        .id(id)
                        .codePointCount(codePointCount)
                        .seqId(seqId++)
                        .documentId(documentId)
                        .status(SegmentStatusEnum.WAITING.name())
                        .segment(segment.text())
                        .createTime(LocalDateTime.now())
                        .updateTime(LocalDateTime.now())
                        .vectorId(UUID.randomUUID()
                                .toString())
                        .title(segment.metadata()
                                .getString("title"))
                        .build();
                kwSegments.add(kwSegment);
            }
            iKwSegmentService.saveBatch(kwSegments);
            int segmentCount = kwSegments.size();
            int codePointCount = kwSegments.stream()
                    .mapToInt(KwSegment::getCodePointCount)
                    .sum();
            EmbeddingStoreIngestorStrategy embeddingStoreIngestorStrategy = EmbeddingStoreIngestorStrategyFactory.ingestorStrategy(EmbeddingModelTypeEnum.LOCAL_BGE_SMALL_ZH_V15_EMBEDDING_MODEL);
            String collectionName = "kw_segment";
            iKwDocumentService.save(KwDocument.builder()
                    .id(documentId)
                    .childId(childId)
                    .status(TaskStatusEnum.PENDING.name())
                    .fileName(fileName)
                    .docType(DocTypeEnum.COMMON.name())
                    .codePointCount(codePointCount)
                    .segmentCount(segmentCount)
                    .createTime(LocalDateTime.now())
                    .updateTime(LocalDateTime.now())
                    .build());
            TransactionSynchronizationManager.registerSynchronization(new TransactionSynchronization() {
                @Override
                public void afterCommit() {
                    boolean enqueue = documentIngestQueue.enqueue(new DocumentIngestTask(
                            documentId,
                            fileName,
                            collectionName,
                            embeddingStoreIngestorStrategy
                    ));
                    if (!enqueue) {
                        log.warn("队列已满，common-任务被丢弃: {}", fileName);
                        iKwSegmentService.updateById(KwSegment.builder()
                                .status(TaskStatusEnum.FAILED.name())
                                .build()
                        );
                    }
                }
            });

        } catch (IOException e) {
            log.error("解析文件失败: {}", fileName, e);
            throw new SysException("文件解析失败");
        }
        return documentId;
    }

    /**
     * 创建文档-对答
     *
     * @param file 文件
     * @return Long
     */
    @Transactional(rollbackFor = Exception.class)
    public Long createReply(MultipartFile file) {
        Long childId = SecurityContextHolder.get()
                .getChildId();
        Long documentId = IdWorker.getId();
        String fileName = file.getOriginalFilename();
        try {
            FaqExcelDataListener faqExcelDataListener = new FaqExcelDataListener(
                    iKwSegmentService,
                    iKwSegmentQuestionService,
                    documentId,
                    childId
            );
            FastExcel.read(file.getInputStream(), FaqData.class, faqExcelDataListener)
                    .sheet()
                    .doRead();
            iKwDocumentService.save(KwDocument.builder()
                    .id(documentId)
                    .childId(childId)
                    .status(TaskStatusEnum.PENDING.name())
                    .fileName(fileName)
                    .codePointCount(faqExcelDataListener.getCodePointCount())
                    .segmentCount(faqExcelDataListener.getSegmentCount())
                    .createTime(LocalDateTime.now())
                    .docType(DocTypeEnum.REPLY.name())
                    .updateTime(LocalDateTime.now())
                    .build());
            String collectionName = "kw_segment";
            EmbeddingStoreIngestorStrategy embeddingStoreIngestorStrategy = EmbeddingStoreIngestorStrategyFactory.ingestorStrategy(EmbeddingModelTypeEnum.LOCAL_BGE_SMALL_ZH_V15_EMBEDDING_MODEL);
            TransactionSynchronizationManager.registerSynchronization(new TransactionSynchronization() {
                @Override
                public void afterCommit() {
                    boolean enqueue = documentIngestQueue.enqueue(new DocumentIngestTask(
                            documentId,
                            fileName,
                            collectionName,
                            embeddingStoreIngestorStrategy
                    ));
                    if (!enqueue) {
                        log.warn("队列已满，reply-任务被丢弃: {}", fileName);
                        iKwSegmentService.updateById(KwSegment.builder()
                                .status(TaskStatusEnum.FAILED.name())
                                .build()
                        );
                    }

                }
            });
        } catch (IOException e) {
            log.error("读取reply 文件失败", e);
            throw new SysException("读取reply 文件失败");
        }
        return documentId;
    }

    /**
     * 重新索引
     *
     * @param documentId 文档id
     * @return Long
     */
    @Transactional(rollbackFor = Exception.class)
    public Long reIndex(Long documentId) {
        Long childId = SecurityContextHolder.get()
                .getChildId();
        KwDocument kwDocument = iKwDocumentService.getOne(new LambdaQueryWrapper<KwDocument>()
                .eq(KwDocument::getChildId, childId)
                .eq(KwDocument::getId, documentId));
        if (Objects.isNull(kwDocument)) {
            throw new BizException("文档不存在");
        }
        String fileName = kwDocument.getFileName();
        EmbeddingStoreIngestorStrategy embeddingStoreIngestorStrategy = EmbeddingStoreIngestorStrategyFactory.ingestorStrategy(EmbeddingModelTypeEnum.LOCAL_BGE_SMALL_ZH_V15_EMBEDDING_MODEL);
        String collectionName = "kw_segment";
        TransactionSynchronizationManager.registerSynchronization(new TransactionSynchronization() {
            @Override
            public void afterCommit() {
                boolean enqueue = documentIngestQueue.enqueue(new DocumentIngestTask(
                        documentId,
                        fileName,
                        collectionName,
                        embeddingStoreIngestorStrategy
                ));
                if (!enqueue) {
                    throw new RuntimeException("队列已满，请稍候再试: " + fileName);
                }
            }
        });

        return documentId;
    }

}
